import os,cv2
import numpy as np
from math import ceil

from dataset.data_aug import *
from dataset.dataset import collate_fn
import time
import torch
import torch.utils.data as torchdata
import torch.utils.data as data
from torch.autograd import Variable
# from torchvision import transforms

# from classify3.resnet import *
from models.MobileNet import MobileNet

class dataset_pred(data.Dataset):
    def __init__(self, imgs,label, transforms=None):
        self.imgs = imgs
        self.labels = label
        self.transforms = transforms

    def __len__(self):
        return len(self.imgs)

    def __getitem__(self, item):
        # img_path = self.paths[item]
        img =self.imgs[0]
        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
        if self.transforms is not None:
            img = self.transforms(img)
        label = self.labels[item]
        return torch.from_numpy(img).float(), label

def class5_init(resume):


    test_transforms= Compose([
            ExpandBorder(size=(272,272),resize=True),
            Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    # resnet
    # model =resnet50(pretrained=True)
    # model.avgpool = torch.nn.AdaptiveAvgPool2d(output_size=1)
    # model.fc = torch.nn.Linear(model.fc.in_features,4)

    #mobilenet
    model = MobileNet(5, alpha=0.5)
    # model.fc = torch.nn.Conv2d(in_channels=512, out_channels=5, kernel_size=1, padding=0)
    model.fc = torch.nn.Linear(model.fc.in_features, 5)

    model.load_state_dict(torch.load(resume),strict=True)
    model = model.cuda()
    model.eval()
    return model,test_transforms

def class5_forward(testimage,model,test_transforms):

    ## 4  其他
    # 3 竖直
    # 2水平倾斜
    # 1 亮
    # 0 灭
    start =time.time()

    input = test_transforms(cv2.cvtColor(testimage[0],cv2.COLOR_BGR2RGB))
    input = input[np.newaxis, :]
    input = torch.FloatTensor(input)
    input = input.to(torch.device("cuda"))
    output = model(input)

    if isinstance(output, list):
        output = (output[0] + output[1]) / 2

    # print(output)
    props, preds = torch.max(output, 1)

    test_preds = preds.data.cpu().numpy()[0]

    # print('predict used {} s'.format(time.time() - start))
    return test_preds


if __name__ == '__main__':

    #########   set the GPU   ###########
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    #########  the input path   ##########
    import glob
    import tqdm

    bar = tqdm.tqdm()
    resume = '/home/heils_lhl/Heils/youi-meter-p/classify3/models/weights-11-246-[1.0000].pth'
    model, test_transforms = class5_init(resume)

    for item in glob.iglob('/media/heils_lhl/data/youirobotData/light/*/*.jpg'):
        # item = '/media/heils_lhl/data/youirobotData/light/2/1577264278.065616_ .jpg'


        img_ori = cv2.imread(item)
        # img = cv2.cvtColor(img_ori,cv2.COLOR_BGR2RGB)
        img = img_ori
        testimage = [img]
        result = class5_forward(testimage,model,test_transforms)
        # print(result[0][0])
        # print(str(result[0][0]))

        gt = int(item.split('/')[-2])
        prd = int(result)
        # print(result,gt,item)
        if gt != prd:
            print(result, gt, item)

            targretfolder = '/media/heils_lhl/data350/validict4/'+str(gt)+'/'+str(result)+'/'
            if not os.path.exists(targretfolder):
                os.makedirs(targretfolder)
            filename = item.split('/')[-1]
            cv2.imwrite(targretfolder+filename,img_ori)
        cv2.imshow('wrong', img_ori)
        cv2.waitKey(0)
        bar.update(1)




        ##########   forward     ###########
        # test_size = ceil(len(data_set['test']) / data_loader['test'].batch_size)



